import open3d as o3d
import numpy as np
from collections import deque
import math


class RegionGrow:

    def __init__(self,
                 cloud,
                 min_pts_per_cluster=1,  # 每个聚类的最小点数
                 max_pts_per_cluster=np.inf,  # 每个聚类的最大点数
                 theta_threshold=30,  # 法向量夹角阈值
                 curvature_threshold=0.05,  # 曲率阈值 0.5:有些曲面残缺不全
                 neighbour_number=30):  # 邻域搜索点数

        """

        :param cloud: 点云数据
        :param min_pts_per_cluster: 每个聚类的最小点数
        :param max_pts_per_cluster: 每个聚类的最大点数
        :param theta_threshold: 法向量夹角阈值（角度）
        :param curvature_threshold: 曲率阈值
        :param neighbour_number: 邻域搜索点数
        """

        self.pcd = cloud
        self.pts_num = len(cloud.points)  # 点云中的点个数
        self.min_pts_per_cluster = min_pts_per_cluster
        self.max_pts_per_cluster = max_pts_per_cluster
        self.theta_threshold = np.deg2rad(theta_threshold)
        self.curvature_threshold = curvature_threshold
        self.neighbour_number = neighbour_number
        self.point_labels = None  # 保存点标签的数组
        self.num_pts_in_segment = []  # 保存每个区域内点个数的list
        self.clusters = []  # 保存每个区域内点云原始索引的双层list
        self.cure = None  # 存储每个点曲率的容器

    # -------------------------------------参数准备--------------------------------------
    def prepare_for_segment(self):
        """
        区域生长前的准备
        :return: 是否可以进行区域生长
        """
        point_nums = len(self.pcd.points)  # 点云个数
        normal_nums = len(self.pcd.normals)  # 点云法向量
        # 判断点云是否为空
        if not point_nums:
            return False
        # 判断查找的邻域点个数是否为0
        if self.neighbour_number == 0:
            return False
        # 点云需要包含法向量信息
        if point_nums != normal_nums:
            self.pcd.estimate_normals(o3d.geometry.KDTreeSearchParamKNN(self.neighbour_number))

        return True

    # -----------------------------------判断邻域点所属分类-----------------------------------
    def validate_points(self, point, nebor):
        """
        判断邻域点是否属于该区域，是否可以作为新的种子点继续生长
        :param point: 当前种子点
        :param nebor: 邻域点
        :return: 是否属于该区域，是否是新的种子点
        """
        is_seed = True  # 初始化种子点标记
        cosine_threshold = np.cos(self.theta_threshold)  # 法向量夹角（平滑）阈值
        curr_seed_normal = self.pcd.normals[point]  # 当前种子点的法向量
        seed_nebor_normal = self.pcd.normals[nebor]  # 种子点邻域点的法向量
        dot_normal = np.fabs(np.dot(seed_nebor_normal, curr_seed_normal))  # 法向量夹角的余弦值

        # 如果法线夹角余弦小于平滑阈值（夹角越大，余弦越小）
        if dot_normal < cosine_threshold:
            # 邻域点为不属于该区域
            return False, is_seed

        # 如果邻域点的曲率大于曲率阈值
        if self.cure[nebor] > self.curvature_threshold:
        # if math.fabs(self.cure[nebor] - self.cure[point]) > self.curvature_threshold:
            # 邻域点属于该区域，但不是种子点
            is_seed = False

        return True, is_seed

    # ----------------------------------对点附上分类标签----------------------------------
    def label_for_points(self, initial_seed, segment_index, tree):
        """
        根据种子点对邻域点进行分类
        :param initial_seed: 原始点云中曲率最小的点作为初始的种子点
        :param segment_index: 区域id
        :param tree: 原始点云的KDtree
        :return: 该区域中点云的个数
        """
        seeds = deque([initial_seed])  # 转成双向队列，类似于list的容器，可以快速的在队列头部和尾部添加、删除元素
        self.point_labels[initial_seed] = segment_index  # 为初始种子点标记当前区域id
        num_pts_in_segment = 1  # 统计该区域中的点个数

        while len(seeds):
            # 当种子点个数不为0时，执行分类代码
            curr_seed = seeds[0]  # 选择当前种子点
            seeds.popleft()  # 在种子点集合中删除当前种子点（删除最左边的元素）
            # 搜索邻域内的点
            [_, seed_neighbours, _] = tree.search_knn_vector_3d(self.pcd.points[curr_seed], self.neighbour_number)

            for i_nebor in seed_neighbours:
                # 遍历当前种子点的邻域点
                if self.point_labels[i_nebor] != -1:
                    # 如果邻域点已被标记，则跳过该邻域点
                    continue

                belongs_to_segment, is_seed = self.validate_points(curr_seed, i_nebor)  # 判断邻域点是否属于该区域，是否是新的种子点

                if not belongs_to_segment:
                    # 如果邻域点不属于该区域，则跳过该点
                    continue

                self.point_labels[i_nebor] = segment_index  # 为邻域点标记当前区域id
                num_pts_in_segment += 1  # 更新区域内点个数

                if is_seed:
                    # 如果邻域点是新的种子点，则添加到种子点列表中
                    seeds.append(i_nebor)

        if self.min_pts_per_cluster <= num_pts_in_segment <= self.max_pts_per_cluster:
            segment_indexes = np.where(self.point_labels == segment_index)[0]  # 提取当前聚类内的所有点云的原始索引
            self.clusters.append(segment_indexes)  # 将该区域内的点云的原始索引保存到聚类合集中
        # else:
        #     # self.point_labels[segment_indexes] = -1
        #     pass
        return num_pts_in_segment

    # ------------------------------------区域生长过程------------------------------------
    def region_growing_process(self):
        """
        区域生长代码
        :return: None
        """
        self.point_labels = -np.ones(self.pts_num)  # 初始化点云标签，-1表示未标记
        self.pcd.estimate_covariances(o3d.geometry.KDTreeSearchParamKNN(self.neighbour_number))  # 计算点云的协方差矩阵
        cov_mat = self.pcd.covariances  # 获取点云的协方差矩阵

        # 计算并统计每个点的曲率
        point_curvatures = np.zeros((self.pts_num, 2))  # 初始化存储每个点曲率和原始索引的容器
        eigenvalues, _ = np.linalg.eig(cov_mat)  # SVD分解求特征值
        point_curvatures[:, 0] = np.min(eigenvalues, axis=1) / np.sum(eigenvalues, axis=1)  # 第1列为曲率
        point_curvatures[:, 1] = np.arange(self.pts_num)  # 第2列为点云的原始索引
        self.cure = point_curvatures[:, 0]  # 存储每个点曲率的容器

        # 按照曲率大小进行排序
        temp_cure = np.argsort(point_curvatures[:, 0])  # 排列后的顺序
        point_curvatures = point_curvatures[temp_cure, :]  # 按从小到大顺序排列后的曲率

        # 选取曲率最小的点作为初始种子点
        seed_counter = 0  # 初始种子点在排序后的索引
        seed = int(point_curvatures[seed_counter, 1])  # 初始种子在原始点云中的索引

        segmented_pts_num = 0  # 已聚类的点云个数
        segment_id = 0  # 区域的id

        kdtree = o3d.geometry.KDTreeFlann(self.pcd)  # 创建KDtree

        while segmented_pts_num < self.pts_num:
            # 点云还未被全部聚类之前，执行区域生长代码
            pts_in_segment = self.label_for_points(seed, segment_id, kdtree)  # 根据当前种子点进行区域生长，并返回该区域生长结束后的点云个数

            segmented_pts_num += pts_in_segment  # 更新已聚类的点云个数
            # self.num_pts_in_segment.append(pts_in_segment)  # 保存该区域内点个数
            segment_id += 1  # 区域id+1，开始下一区域的生长

            # 寻找下一个种子点进行下一区域的生长聚类
            for i_seed in range(seed_counter + 1, self.pts_num):
                # 遍历按曲率排序后的当前种子点之后的点云
                index = int(point_curvatures[i_seed, 1])
                if self.point_labels[index] == -1:
                    # 若该点还未被标记，则作为下一区域的初始种子点
                    seed = index  # 新的种子点在原始点云中的索引
                    seed_counter = i_seed  # 新的种子点在重新排序后的索引
                    break
        print(f"原始聚类个数为 {segment_id}")

    # ----------------------------------根据标签进行分类-----------------------------------
    def region_growing_clusters(self):
        """
        统计每个区域内点云个数，并分别保存区域内点云的原始索引
        :return: None
        """
        num_of_segments = len(self.num_pts_in_segment)  # 生长出的区域个数

        # 初始化聚类数组
        for segment_id in range(num_of_segments):
            # 遍历所有区域
            segment_index = np.where(self.point_labels == segment_id)[0]  # 提取区域内所有点云的原始索引
            if self.min_pts_per_cluster <= self.num_pts_in_segment[segment_id] <= self.max_pts_per_cluster:
                # 区域内点云个数符合要求
                self.clusters[segment_id] = list(segment_index)  # 将每个区域内点云的原始索引保存成list并添加到聚类合集中
            else:
                # 区域内点云个数不符合要求
                self.point_labels[segment_index] = -1  # 将该区域内的点云标记为未分类

    # ----------------------------------执行区域生长算法-----------------------------------
    def extract(self):
        """
        调用函数执行区域生长
        :return: 筛选后的区域点云索引
        """
        if not self.prepare_for_segment():
            # 判断点云数据是否存在问题
            print("区域生长算法预处理失败！")
            return

        self.region_growing_process()  # 区域生长代码

        # self.region_growing_clusters()

        return self.clusters, self.cure
